import os

import cv2
from tqdm import tqdm

from option import parse_args


def unaligned_videos(original_path, manipulated_path, compression_version, func, template: str):
    real_fake_pairs = []
    for manipulation_type in os.listdir(manipulated_path):
        manipulation_type_path = os.path.join(manipulated_path, manipulation_type)
        manipulation_videos_path = os.path.join(manipulation_type_path, compression_version, 'videos')
        manipulation_videos = os.listdir(manipulation_videos_path)
        manipulation_videos.sort()
        for video_name in manipulation_videos:
            video_path_split = video_name.split('_')
            assert len(video_path_split) == 2
            original = '{}.mp4'.format(video_path_split[0])
            real_video_path = os.path.join(original_path, original)
            fake_video_path = os.path.join(manipulation_videos_path, video_name)
            real_fake_pairs.append((real_video_path, fake_video_path))
    print(len(real_fake_pairs))

    for real_video_path, fake_video_path in tqdm(real_fake_pairs):
        capture = cv2.VideoCapture(real_video_path)
        real_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        capture = cv2.VideoCapture(fake_video_path)
        fake_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if fake_video_frames_num != real_video_frames_num:
            func(template.format(fake_video_path, fake_video_frames_num - real_video_frames_num))


def main():
    args = parse_args()
    assert os.path.exists(args.root_dir)
    original_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', args.compression_version, 'videos')
    manipulated_path = os.path.join(args.root_dir, 'manipulated_sequences')
    args.save = True
    if args.save:
        with open(f'unaligned_videos_statistics_{args.compression_version}.dat', 'w') as f:
            unaligned_videos(original_path, manipulated_path, args.compression_version, f.write, '{},{}\n')
    else:
        unaligned_videos(original_path, manipulated_path, args.compression_version, print, '{},{}')


if __name__ == '__main__':
    main()
